#!/usr/bin/env python

from __future__ import print_function, division, unicode_literals

__doc__ = """filter txinfo files by given criteria"""

import argparse
import os
import sys

import uta.formats.txinfo as ufti


def parse_args(argv):
    ap = argparse.ArgumentParser(
        description=__doc__,
    )
    ap.add_argument(
        "FILE"
    )
    ap.add_argument(
        "--genes-from", "-G",
    )
    ap.add_argument(
        "--transcripts-from", "-T",
    )

    opts = ap.parse_args(argv)
    return opts


if __name__ == "__main__":
    opts = parse_args(sys.argv[1:])

    filter_genes = set()
    if opts.genes_from:
        filter_genes = set([l.strip()
                            for l in open(opts.genes_from, "r").readlines()])

    filter_acs = set()
    if opts.transcripts_from:
        filter_acs = set([l.strip()
                          for l in open(opts.transcripts_from, "r").readlines()])

    tir = ufti.TxInfoReader(
        sys.stdin if opts.FILE == "-" else open(opts.FILE, "r"))
    tiw = ufti.TxInfoWriter(sys.stdout)

    for ti in tir:
        if ti.hgnc in filter_genes or ti.ac in filter_acs:
            tiw.write(ti)
